from weaviate.collections.classes.grpc import MetadataQuery

from db.database import db

from typing import List

from langchain_core.callbacks import CallbackManagerForRetrieverRun, AsyncCallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

from utils.tools import llm_split, create_vector, config


class WeaviateHybridRetriever(BaseRetriever):
    """Weaviate 混合检索，传入空格分割的query和query向量做查询
    """

    top_k: int
    """Number of top results to return"""

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Sync implementations for retriever."""
        vector = create_vector(query)
        query = llm_split(query)

        matching_documents = []
        client = db.get_client()
        reviews = client.collections.get(config["weaviate"]["search_collection"]["name"])
        query_result = reviews.query.hybrid(
            query=query,
            vector=vector,
            limit=self.top_k,
            # alpha=0.75,
            # certainty=0.90,
            # distance=0.25,
            return_metadata=MetadataQuery(distance=True, certainty=True),
        )
        for doc in query_result.objects:
            tmp_doc = Document(page_content="", metadata={})
            for key in doc.properties:
                if key == 'text':
                    tmp_doc.page_content = doc.properties[key]
                else:
                    tmp_doc.metadata[key] = doc.properties[key]
            matching_documents.append(tmp_doc)

        return matching_documents


    async def _aget_relevant_documents(
        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Asynchronously get documents relevant to a query.

        Args:
            query: String to find relevant documents for
            run_manager: The callbacks handler to use

        Returns:
            List of relevant documents
        """